from SPARQLWrapper import SPARQLWrapper, JSON


class BlanketRetriever(object):
    sparqlConn = SPARQLWrapper("http://dbpedia.org/sparql");
    sparqlConn.setReturnFormat(JSON);
    def __init__(self):
        self.__blanket = set([]);
        self.__category = "";
    
    def grabCategoryBlanket(self, category):
        self.__category = category;
        # Grab both the supercats and subcats...
        super2 = self.__runSuper(category);
        (super, sub) = self.__grabSuperSub();
        # Grab the sibling categories...
        sib = self.__grabSibling(sub);
        self.__blanket = self.__blanket.union(super).union(sub).union(sib).union(super2);
        file3 = open("/Users/Hikari/Documents/workspace/WikiDataSearch/WikiCategories/"+category+".categories","a");
        for cat in self.__blanket:
            file3.write(cat+"\n");
        file3.close();
        print (str(len(self.__blanket)) + " total categories in this blanket.");
        return (self.__blanket, sub);
    
    def __grabSuperSub(self):
        sub = set([]);
        super = set([]);
        (supers, subs) = self.__runSuperSub();
        subCount = 0;
        superCount = 0;
        for subcategory in subs:
            sub.add(subcategory);
            subCount += 1;
        for supercategory in supers:
            super.add(supercategory);
            superCount += 1;
        # Subcategory printout
        if(subCount == 0): print("No subcategories found.");
        elif(subCount == 1): print("1 subcategory found.");
        else: print(str(subCount) + " subcategories found.");
        # Supercategory printout
        if(superCount == 0): print("No supercategories found.");
        elif(superCount == 1): print("1 supercategory found.");
        else: print(str(superCount) + " supercategories found");
        return (super, sub);
    
    def __grabSibling(self, subs):
        sib = set([]);
        for category in subs:
            sib = self.__runSuper(category[37:]);
        if(len(sib) == 0): print("No sibling categories found.");
        elif(len(sib) == 1): print("1 sibling category found.");
        else: print(str(len(sib)) + " sibling categories found.");
        return sib;
    
    def __runSuper(self, cat):
        sib = set([]);
        count = 0;
        query = self.__buildSuperQuery(cat);
        BlanketRetriever.sparqlConn.setQuery(query);
        resultSet = BlanketRetriever.sparqlConn.query().convert();
        for result in resultSet["results"]["bindings"]:
            sib.add(result["broader"]["value"]);
            count += 1;
        if(len(sib) == 0): print("No supers for "+cat+" found.");
        elif(len(sib) == 1): print("1 super for "+cat+" found.");
        else: print(str(len(sib)) + " supers for "+cat+" found.");
        return sib;
    
    def __runSuperSub(self):
        subcat = [];
        supercat = [];
        query = self.__buildSuperSubQuery();
        BlanketRetriever.sparqlConn.setQuery(query);
        resultSet = BlanketRetriever.sparqlConn.query().convert();
        for result in resultSet["results"]["bindings"]:
            if(len(result["o"]["value"]) > 36 ):
                if(result["o"]["value"][37:] == self.__category): #We will be grabbing a subclass...
                    if(len(result["s"]["value"]) > 36):
                        if(result["s"]["value"][36] == ":"):
                            subcat.append(result["s"]["value"]);
                else: # we are going to grab a superclass...
                    if(result["o"]["value"][36] == ":"): # Then it really is a superclass...
                        supercat.append(result["o"]["value"]);
        return (supercat, subcat);
    
    
    def __buildSuperQuery(self, cat= None):
        if (cat == None): cat = self.__category;
        query = ("PREFIX rdfs: <http://www.w3.org/2000/01/rdf-schema#> "
                 "SELECT ?broader "
                 "WHERE { <http://dbpedia.org/resource/Category:"+ cat +"> skos:broader ?broader }");
        return query;
    def __buildSuperSubQuery(self):
        query = ("DESCRIBE <http://dbpedia.org/resource/Category:"+ self.__category +">");
        return query;